#include "olsfitter.h"
#include <algorithm>
#include <QDebug>
#include <cmath>
#include <iostream>
#include <functional>

OLSFitter::OLSFitter() = default;


Vecd gaussSolve(Matd &A, Vecd &B) {
    auto n = A.rows();
    auto m = n + 1;
    Matd U(n, m);
    U.block(0, 0, n, n) = A;
    U.col(n) = B;
    Matd tmp;

//    化阶梯形
    for (size_t k = 0; k < n - 1; ++k) {
//      U[k+1:n-1,k:m-1] -= U[k+1:n-1,k] * U[k,k:m-1]/U(k,k)   *为这里为矩阵乘法
        U.block(k + 1, k, n - 1 - k, m - k) -= U.col(k).tail(n - 1 - k) / U(k, k) * U.row(k).tail(m - k);

    }
    //回代
    for (int i = static_cast<int>(n) - 1; i >= 0; --i) {
//      U[i,i:m-1] /= U[i,i]
//      U.row(i).tail(cnt(i,m-1)) /= U(i,i);
        U.row(i).tail(m - i) /= U(i, i);
//      U[0:i-1,i:m-1] -= U[0:i-1,i]*U[i,i:m-1]
//      U.block(0,i,cnt(0,i-1),cnt(i,m-1)) -= U.col(i).head(cnt(0,i-1))*U.row(i).tail(cnt(i,m-1));
        U.block(0, i, i, m - i) -= U.col(i).head(i) * U.row(i).tail(m - i);

    }
    return U.col(n);
}

void OLSFitter::setExp() {
    std::vector<std::function<double(double)>> g(2);
    g[0] = [](double x) { return 1.0; };
    g[1] = [](double x) { return x; };
    auto eval = static_cast<double (*)(double)>(std::exp);
    auto processY = [](double x){
        //需要处理确诊人数为0的情况，因为log函数的自变量要是正整数
        //人数很小的时候没有多大意义，所以直接把它变成1
        x = x<1?1:x;
        return std::log(x);
    };
    setFormulas(g, processY, eval);
}

void OLSFitter::setPoly(int n) {
    if (n < 1) {
        return;
    }

    std::vector<std::function<double(double)>> g(n + 1);
    g[0] = [](double x) { return 1.0; };

    for (size_t i = 1; i <= static_cast<size_t>(n); ++i) {
        g[i] = [=](double x) { return std::pow(x, i); };
    }
    auto eval = [](double x) { return x; };
    auto processY = [](double x) { return x; };
    setFormulas(g, processY, eval);
}

bool OLSFitter::setFormulas(const std::vector<std::function<double(double)>> &g,
                            const std::function<double(double)> &processY, const std::function<double(double)> &eval) {
    for (const auto &func: g) {
        if (!func) {
            return false;
        }
    }
    if (!eval) {
        return false;
    }
    if (!processY) {
        return false;
    }
    m_g = g;
    m_eval = eval;
    m_processY = processY;
    calcCoefficients();
    return true;

}



bool OLSFitter::setPoints(const Vecd &x, const Vecd &y) {
    if (x.size() != y.size() || x.size() == 0 || x.size() < m_g.size()) {
        return false;
    }
    m_x = x;
    m_y = y.unaryExpr(m_processY);
    calcCoefficients();
    return true;
}

void OLSFitter::calcCoefficients() {
    if (m_g.empty()) {
        return;
    }
    if (!m_eval) {
        return;
    }
    if (m_x.size() != m_y.size() || m_x.size() == 0 || m_x.size() < m_g.size()) {
        return;
    }
    auto n_p = m_x.size();
    auto n_g = m_g.size();
    Matd Gt(n_g, n_p);
    for (size_t i = 0; i < n_g; ++i) {
        Gt.row(i) = m_x.unaryExpr(m_g[i]);
    }


    Matd A = Gt * Gt.transpose();
    Vecd B = Gt * m_y.matrix();
    std::cout<<"A"<<std::endl;
    std::cout<<A<<std::endl;
    std::cout<<"B"<<std::endl;
    std::cout<<B<<std::endl;

    m_a = gaussSolve(A, B);
    std::cout<<"系数："<<m_a<<std::endl;


}

double OLSFitter::operator()(double x) {
    double ret = 0;
    for (size_t i = 0; i < m_g.size(); ++i) {
        ret += m_a[i] * m_g.at(i)(x);
    }
    return m_eval(ret);
}
